// SSH WebSocket服务器示例
// 这是一个Node.js后端服务器示例，用于处理前端终端的SSH连接请求
// 使用方法：node ssh-server-example.js

const WebSocket = require('ws');
const { Client } = require('ssh2');
const http = require('http');
const url = require('url');
const jwt = require('jsonwebtoken'); // 需要安装: npm install jsonwebtoken

// JWT密钥（生产环境应该从环境变量获取）
const JWT_SECRET = process.env.JWT_SECRET || 'your-secret-key';

// 创建HTTP服务器
const server = http.createServer();

// Token验证函数
function verifyToken(token) {
  try {
    // 这里可以根据实际情况验证token
    // 方案1: JWT验证
    const decoded = jwt.verify(token, JWT_SECRET);
    return { valid: true, user: decoded };
    
    // 方案2: 简单字符串验证（仅用于演示）
    // if (token === 'valid-token') {
    //   return { valid: true, user: { id: 1, username: 'admin' } };
    // }
    
    // 方案3: 数据库验证（需要连接数据库）
    // const user = await getUserByToken(token);
    // return user ? { valid: true, user } : { valid: false };
    
  } catch (error) {
    console.error('Token验证失败:', error);
    return { valid: false, error: error.message };
  }
}

// 创建WebSocket服务器
const wss = new WebSocket.Server({ 
  server,
  path: '/terminal/ssh',
  // 连接验证回调
  verifyClient: (info) => {
    // 方案1: 从URL参数获取token
    const query = url.parse(info.req.url, true).query;
    const token = query.token;
    
    if (!token) {
      console.log('WebSocket连接被拒绝: 缺少token');
      return false;
    }
    
    const verification = verifyToken(token);
    if (!verification.valid) {
      console.log('WebSocket连接被拒绝: token无效');
      return false;
    }
    
    // 将用户信息附加到请求对象上
    info.req.user = verification.user;
    return true;
  }
});

console.log('SSH WebSocket服务器启动中...');

// 存储SSH连接
const sshConnections = new Map();

wss.on('connection', (ws, req) => {
  console.log('新的WebSocket连接建立');
  console.log('用户信息:', req.user);
  
  let sshClient = null;
  let sshStream = null;
  let isAuthenticated = true; // 通过verifyClient已验证
  
  ws.on('message', (message) => {
    try {
      const data = JSON.parse(message.toString());
      
      switch (data.type) {
        case 'auth':
          // 方案3: 连接后认证
          handleAuth(ws, data);
          break;
          
        case 'connect':
          // 检查认证状态
          if (!isAuthenticated) {
            ws.send(JSON.stringify({
              type: 'error',
              message: '未认证，请先进行身份验证'
            }));
            return;
          }
          
          // 方案4: 在业务消息中验证token
          if (data.token) {
            const verification = verifyToken(data.token);
            if (!verification.valid) {
              ws.send(JSON.stringify({
                type: 'error',
                message: 'Token验证失败: ' + (verification.error || '无效的token')
              }));
              return;
            }
            console.log('业务消息token验证成功:', verification.user);
          }
          
          handleSSHConnect(ws, data);
          break;
          
        case 'input':
          if (sshStream) {
            sshStream.write(data.data);
          }
          break;
          
        case 'resize':
          if (sshStream) {
            sshStream.setWindow(data.rows, data.cols);
          }
          break;
          
        case 'disconnect':
          handleDisconnect();
          break;
      }
    } catch (error) {
      console.error('解析消息失败:', error);
      ws.send(JSON.stringify({
        type: 'error',
        message: '消息格式错误'
      }));
    }
  });
  
  // 处理认证（方案3）
  function handleAuth(ws, data) {
    const verification = verifyToken(data.token);
    if (verification.valid) {
      isAuthenticated = true;
      ws.send(JSON.stringify({
        type: 'auth_success',
        message: '认证成功',
        user: verification.user
      }));
      console.log('用户认证成功:', verification.user);
    } else {
      isAuthenticated = false;
      ws.send(JSON.stringify({
        type: 'auth_failed',
        message: 'Token验证失败: ' + (verification.error || '无效的token')
      }));
      console.log('用户认证失败');
    }
  }
  
  // 处理SSH连接
  function handleSSHConnect(ws, data) {
    const { config, terminalSize } = data;
    
    sshClient = new Client();
    
    sshClient.on('ready', () => {
      console.log('SSH客户端连接就绪');
      
      sshClient.shell({
        cols: terminalSize.cols,
        rows: terminalSize.rows,
        term: 'xterm-256color'
      }, (err, stream) => {
        if (err) {
          console.error('创建shell失败:', err);
          ws.send(JSON.stringify({
            type: 'error',
            message: '创建shell失败: ' + err.message
          }));
          return;
        }
        
        sshStream = stream;
        
        // 发送连接成功消息
        ws.send(JSON.stringify({
          type: 'connected'
        }));
        
        // 监听SSH输出
        stream.on('data', (data) => {
          ws.send(JSON.stringify({
            type: 'data',
            data: data.toString()
          }));
        });
        
        // 监听SSH错误
        stream.on('error', (err) => {
          console.error('SSH流错误:', err);
          ws.send(JSON.stringify({
            type: 'error',
            message: 'SSH流错误: ' + err.message
          }));
        });
        
        // 监听SSH关闭
        stream.on('close', () => {
          console.log('SSH流已关闭');
          ws.send(JSON.stringify({
            type: 'disconnected'
          }));
        });
      });
    });
    
    sshClient.on('error', (err) => {
      console.error('SSH连接错误:', err);
      ws.send(JSON.stringify({
        type: 'error',
        message: 'SSH连接失败: ' + err.message
      }));
    });
    
    // 连接SSH服务器
    try {
      sshClient.connect({
        host: config.host,
        port: config.port,
        username: config.username,
        password: config.password,
        readyTimeout: 20000
      });
    } catch (error) {
      console.error('SSH连接异常:', error);
      ws.send(JSON.stringify({
        type: 'error',
        message: 'SSH连接异常: ' + error.message
      }));
    }
  }
  
  // 处理断开连接
  function handleDisconnect() {
    if (sshStream) {
      sshStream.end();
      sshStream = null;
    }
    if (sshClient) {
      sshClient.end();
      sshClient = null;
    }
  }
  
  // WebSocket连接关闭
  ws.on('close', () => {
    console.log('WebSocket连接已关闭');
    handleDisconnect();
  });
  
  // WebSocket错误
  ws.on('error', (error) => {
    console.error('WebSocket错误:', error);
    handleDisconnect();
  });
});

// 启动服务器
const PORT = process.env.PORT || 3000;
server.listen(PORT, () => {
  console.log(`SSH WebSocket服务器运行在端口 ${PORT}`);
  console.log(`WebSocket地址: ws://localhost:${PORT}/ssh`);
});

// 优雅关闭
process.on('SIGINT', () => {
  console.log('\n正在关闭服务器...');
  
  // 关闭所有SSH连接
  sshConnections.forEach((connection) => {
    if (connection.sshClient) {
      connection.sshClient.end();
    }
  });
  
  // 关闭WebSocket服务器
  wss.close(() => {
    console.log('WebSocket服务器已关闭');
    process.exit(0);
  });
});

/*
安装依赖:
npm install ws ssh2

运行服务器:
node ssh-server-example.js

注意事项:
1. 这是一个基础示例，生产环境需要添加更多安全措施
2. 建议添加身份验证和授权机制
3. 可以添加连接池管理多个SSH连接
4. 建议添加日志记录和错误处理
5. 可以添加SSL/TLS支持(wss://)
*/